# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.  
# SPDX-License-Identifier: MPL-2.0

# AUTHORS: Josep Fanals Batllori and Santiago Peñate Vera
# CONTACT:  u1946589@campus.udg.edu and santiago.penate.vera@gmail.com
# thanks to Llorenç Fanals Batllori for his help at coding
import pandas as pd
import numpy as np
import numba as nb
import time
from warnings import warn
from scipy.sparse import csc_matrix, coo_matrix
from scipy.sparse import hstack as hs, vstack as vs
from scipy.sparse.linalg import spsolve, factorized

from VeraGridEngine.DataStructures.numerical_circuit import NumericalCircuit
from VeraGridEngine.Simulations.PowerFlow.power_flow_results import NumericPowerFlowResults
import VeraGridEngine.Simulations.PowerFlow.NumericalMethods.common_functions as cf
from VeraGridEngine.basic_structures import Logger, CscMat, CxVec, IntVec


def epsilon(Sn, n, E):
    """
    Fast recursive Wynn's epsilon algorithm from:
        NONLINEAR SEQUENCE TRANSFORMATIONS FOR THE ACCELERATION OF CONVERGENCE
        AND THE SUMMATION OF DIVERGENT SERIES

        by Ernst Joachim Weniger
    Args:
        Sn: sum of coefficients
        n: order
        E: Coefficients structure copy that is modified in this algorithm

    Returns:

    """
    Zero = complex(0)
    One = complex(1)
    Tiny = np.finfo(complex).tiny  # np.finfo(complex).min
    Huge = np.finfo(complex).max

    E[n] = Sn

    if n == 0:
        estim = Sn
    else:
        AUX2 = Zero

        for j in range(n, 0, -1):  # range from n to 1 (both included)
            AUX1 = AUX2
            AUX2 = E[j - 1]
            DIFF = E[j] - AUX2

            if abs(DIFF) <= Tiny:
                E[j - 1] = Huge
            else:
                if DIFF == 0:
                    DIFF = Tiny
                E[j - 1] = AUX1 + One / DIFF

        if np.mod(n, 2) == 0:
            estim = E[0]
        else:
            estim = E[1]

    return estim, E


# @nb.njit("(c16[:])(i8, c16[:, :], f8)")
def pade4all(order, coeff_mat, s=1.0):
    """
    Computes the "order" Padè approximant of the coefficients at the approximation point s

    Arguments:
        coeff_mat: coefficient matrix (order, buses)
        order:  order of the series
        s: point of approximation (at 1 you get the voltage)

    Returns:
        Padè approximation at s for all the series
    """
    nbus = coeff_mat.shape[1]

    # complex_type = nb.complex128
    complex_type = np.complex128

    voltages = np.zeros(nbus, dtype=complex_type)

    nn = int(order / 2)
    L = nn
    M = nn

    for d in range(nbus):

        # formation of the linear system right hand side
        rhs = coeff_mat[L + 1:L + M + 1, d]

        # formation of the coefficients matrix
        C = np.zeros((L, M), dtype=complex_type)
        for i in range(L):
            k = i + 1
            C[i, :] = coeff_mat[L - M + k:L + k, d]

        # Obtaining of the b coefficients for orders greater than 0
        b = np.zeros(rhs.shape[0] + 1, dtype=complex_type)
        x = np.linalg.solve(C, -rhs)  # bn to b1
        b[0] = 1
        b[1:] = x[::-1]

        # Obtaining of the coefficients 'a'
        a = np.zeros(L + 1, dtype=complex_type)
        a[0] = coeff_mat[0, d]
        for i in range(L):
            val = complex_type(0)
            k = i + 1
            for j in range(k + 1):
                val += coeff_mat[k - j, d] * b[j]
            a[i + 1] = val

        # evaluation of the function for the value 's'
        p = complex_type(0)
        q = complex_type(0)
        for i in range(L + 1):
            p += a[i] * s ** i
            q += b[i] * s ** i

        voltages[d] = p / q

    return voltages


# @nb.njit("(c16[:])(c16[:, :], c16[:, :], i8, c16[:])")
@nb.njit(cache=True)
def sigma_function(coeff_matU, coeff_matX, order, V_slack):
    """

    :param coeff_matU: array with voltage coefficients
    :param coeff_matX: array with inverse conjugated voltage coefficients
    :param order: should be prof - 1
    :param V_slack: slack bus voltage vector. Must contain only 1 slack bus
    :return: sigma complex value
    """
    if len(V_slack) > 1:
        print('Sigma values may not be correct')
    V0 = V_slack[0]
    coeff_matU = coeff_matU / V0
    coeff_matX = coeff_matX / V0
    nbus = coeff_matU.shape[1]
    complex_type = nb.complex128
    # complex_type = complex
    sigmas = np.zeros(nbus, dtype=complex_type)

    if order % 2 == 0:
        M = int(order / 2) - 1
    else:
        M = int(order / 2)

    for d in range(nbus):
        a = coeff_matU[1:2 * M + 2, d]
        b = coeff_matX[0:2 * M + 1, d]
        C = np.zeros((2 * M + 1, 2 * M + 1), dtype=complex_type)

        for i in range(2 * M + 1):
            if i < M:
                C[1 + i:, i] = a[:2 * M - i]
            else:
                C[i - M:, i] = - b[:3 * M - i + 1]

        lhs = np.linalg.solve(C, -a)

        sigmas[d] = np.sum(lhs[M:]) / (np.sum(lhs[:M]) + 1)

    return sigmas


# @nb.njit("(c16[:])(c16[:, :], c16[:, :], i8, i8[:])")
@nb.njit(cache=True)
def conv1_old(A, B, c, indices):
    """
    Performs the convolution of A* and B
    :param A: Coefficients matrix 1 (orders, buses)
    :param B: Coefficients matrix 2 (orders, buses)
    :param c: order of the coefficients
    :param indices: bus indices array
    :return: Array with the convolution for the buses given by "indices"
    """
    suma = np.zeros(len(indices), dtype=nb.complex128)
    for k in range(1, c + 1):
        for i, d in enumerate(indices):
            suma[i] += np.conj(A[k, d]) * B[c - k, d]
    return suma


# @nb.njit("(c16[:])(c16[:, :], c16[:, :], i8)")
@nb.njit(cache=True)
def conv1(A, B, c):
    """
    Performs the convolution of A* and B
    :param A: Coefficients matrix 1 (orders, buses)
    :param B: Coefficients matrix 2 (orders, buses)
    :param c: order of the coefficients
    :return: Array with the convolution for the buses given by "indices"
    """
    suma = np.zeros(A.shape[1], dtype=nb.complex128)
    for k in range(1, c + 1):
        for i in range(A.shape[1]):
            suma[i] += np.conj(A[k, i]) * B[c - k, i]
    return suma


# @nb.njit("(c16[:])(c16[:, :], c16[:, :], i8, i8[:])")
@nb.njit(cache=True)
def conv2(A, B, c, indices):
    """
    Performs the convolution of A and B
    :param A: Coefficients matrix 1 (orders, buses)
    :param B: Coefficients matrix 2 (orders, buses)
    :param c: order of the coefficients
    :param indices: bus indices array
    :return: Array with the convolution for the buses given by "indices"
    """
    suma = np.zeros(len(indices), dtype=nb.complex128)
    for k in range(1, c):
        for i, d in enumerate(indices):
            suma[i] += A[k, d] * B[c - 1 - k, d]
    return suma


# @nb.njit("(c16[:])(c16[:, :], c16[:, :], i8, i8[:])")
@nb.njit(cache=True)
def conv3(A, B, c, indices):
    """
    Performs the convolution of A and B*
    :param A: Coefficients matrix 1 (orders, buses)
    :param B: Coefficients matrix 2 (orders, buses)
    :param c: order of the coefficients
    :param indices: bus indices array
    :return: Array with the convolution for the buses given by "indices"
    """
    suma = np.zeros(len(indices), dtype=nb.complex128)
    for k in range(1, c):
        for i, d in enumerate(indices):
            suma[i] += A[k, d] * np.conj(B[c - k, d])
    return suma


def helm_coefficients_josep(Ybus: CscMat, Yseries: CscMat, V0: CxVec, S0: CxVec, Ysh0: CxVec,
                            pq: IntVec, pv: IntVec, sl: IntVec, no_slack: IntVec,
                            tolerance=1e-6, max_coeff=30, verbose=False, stop_if_too_bad=False,
                            logger: Logger = None):
    """
    Holomorphic Embedding LoadFlow Method as formulated by Josep Fanals Batllori in 2020
    This function just returns the coefficients for further usage in other routines
    :param Ybus: Admittance matrix
    :param Yseries: Admittance matrix of the series elements
    :param V0: vector of specified voltages
    :param S0: vector of specified power
    :param Ysh0: vector of shunt admittances (including the shunts of the Branches)
    :param pq: list of pq nodes
    :param pv: list of pv nodes
    :param sl: list of slack nodes
    :param no_slack: sorted list of pq and pv nodes
    :param tolerance: target error (or tolerance)
    :param max_coeff: maximum number of coefficients
    :param verbose: print intermediate information
    :param stop_if_too_bad: stop when things go wrong
    :param logger: Logger object to store the debug info
    :return: U, X, Q, V, iterations
    """

    n_no_slack = len(no_slack)
    npv = len(pv)
    nsl = len(sl)
    n = Yseries.shape[0]

    # --------------------------- PREPARING IMPLEMENTATION -------------------------------------------------------------
    U = np.zeros((max_coeff + 1, n_no_slack), dtype=complex)  # voltages
    X = np.zeros((max_coeff + 1, n_no_slack), dtype=complex)  # compute X=1/conj(U)
    Q = np.zeros((max_coeff + 1, n_no_slack), dtype=complex)  # unknown reactive powers

    if n < 2:
        # U, X, Q, V, iter_, converged
        return U, X, Q, V0, 0, False

    if verbose:
        logger.add_debug('Yseries', Yseries.toarray())

        df = pd.DataFrame(data=np.c_[Ysh0.imag, S0.real, S0.imag, np.abs(V0)],
                          columns=['Ysh', 'P0', 'Q0', 'V0'])
        logger.add_debug(df.to_string())

    # build the reduced system
    Yred = Yseries[np.ix_(no_slack, no_slack)]  # admittance matrix without slack buses
    Yslack = -Yseries[np.ix_(no_slack, sl)]  # yes, it is the negative of this
    G = Yred.real.copy()  # real parts of Yij
    B = Yred.imag.copy()  # imaginary parts of Yij
    vec_P = S0.real[no_slack]
    vec_Q = S0.imag[no_slack]
    Vslack = V0[sl]
    Ysh = Ysh0[no_slack]
    Vm0 = np.abs(V0[no_slack])
    vec_W = Vm0 * Vm0

    # indices 0 based in the internal scheme
    nsl_counted = np.zeros(n, dtype=int)
    compt = 0
    for i in range(n):
        if i in sl:
            compt += 1
        nsl_counted[i] = compt

    pq_ = pq - nsl_counted[pq]
    pv_ = pv - nsl_counted[pv]
    pqpv_ = np.sort(np.r_[pq_, pv_])

    # .......................CALCULATION OF TERMS [0] ------------------------------------------------------------------

    if nsl > 1:
        U[0, :] = spsolve(Yred, Yslack.sum(axis=1))
    else:
        U[0, :] = spsolve(Yred, Yslack)

    # if U[0, :].real != 0.0 and U[0, :].imag != 0.0:
    X[0, :] = 1 / np.conj(U[0, :])

    # .......................CALCULATION OF TERMS [1] ------------------------------------------------------------------
    valor = np.zeros(n_no_slack, dtype=complex)

    # get the current Injections that appear due to the slack buses reduction
    I_inj_slack = Yslack[pqpv_, :] * Vslack

    valor[pq_] = (I_inj_slack[pq_]
                  - Yslack[pq_].sum(axis=1).A1
                  + (vec_P[pq_] - vec_Q[pq_] * 1j) * X[0, pq_]
                  - U[0, pq_] * Ysh[pq_])

    valor[pv_] = (I_inj_slack[pv_]
                  - Yslack[pv_].sum(axis=1).A1
                  + (vec_P[pv_]) * X[0, pv_]
                  - U[0, pv_] * Ysh[pv_])

    # compose the right-hand side vector
    RHS = np.r_[
        valor.real,
        valor.imag,
        vec_W[pv_] - (U[0, pv_] * U[0, pv_]).real  # vec_W[pv_] - 1.0
    ]

    # Form the system matrix (MAT)
    Upv = U[0, pv_]
    Xpv = X[0, pv_]
    VRE = coo_matrix((2 * Upv.real, (np.arange(npv), pv_)), shape=(npv, n_no_slack)).tocsc()
    VIM = coo_matrix((2 * Upv.imag, (np.arange(npv), pv_)), shape=(npv, n_no_slack)).tocsc()
    XIM = coo_matrix((-Xpv.imag, (pv_, np.arange(npv))), shape=(n_no_slack, npv)).tocsc()
    XRE = coo_matrix((Xpv.real, (pv_, np.arange(npv))), shape=(n_no_slack, npv)).tocsc()
    EMPTY = csc_matrix((npv, npv))

    MAT = vs(
        (hs((G, -B, XIM)),
         hs((B, G, XRE)),
         hs((VRE, VIM, EMPTY))),
        format='csc'
    )

    if verbose:
        logger.add_debug("MAT", MAT.toarray())

    # factorize (only once)
    # MAT_LU = factorized(MAT.tocsc())

    # solve
    try:
        mat_factorized = factorized(MAT)
    except RuntimeError:
        warn("Unable to factorize HELM coefficients matrix :/")
        return U, X, Q, V0, 0, False

    LHS = mat_factorized(RHS)
    # LHS = spsolve(MAT, RHS)

    # update coefficients
    U[1, :] = LHS[:n_no_slack] + 1j * LHS[n_no_slack:2 * n_no_slack]
    Q[0, pv_] = LHS[2 * n_no_slack:]
    X[1, :] = -X[0, :] * np.conj(U[1, :]) / np.conj(U[0, :])

    # .......................CALCULATION OF TERMS [>=2] ----------------------------------------------------------------
    iter_ = 1
    c = 2
    norm_f_prev = 1e20
    converged = False
    V = np.empty(n, dtype=complex)
    V[sl] = V0[sl]
    V[no_slack] = U[:c, :].sum(axis=0)

    a = n_no_slack
    b = 2 * n_no_slack

    while c <= max_coeff and not converged:  # c defines the current depth

        valor[pq_] = (vec_P[pq_] - vec_Q[pq_] * 1j) * X[c - 1, pq_] - U[c - 1, pq_] * Ysh[pq_]
        valor[pv_] = -1j * conv2(X, Q, c, pv_) - U[c - 1, pv_] * Ysh[pv_] + X[c - 1, pv_] * vec_P[pv_]

        RHS = np.r_[
            valor.real,
            valor.imag,
            -conv3(U, U, c, pv_).real
        ]

        # LHS = spsolve(MAT, RHS)
        LHS = mat_factorized(RHS)

        # update voltage coefficients
        U[c, :] = LHS[:a] + 1j * LHS[a:b]

        # update reactive power
        Q[c - 1, pv_] = LHS[b:]

        # update voltage inverse coefficients
        X[c, :] = -conv1(U, X, c) / np.conj(U[0, :])

        # compute power mismatch
        V[no_slack] += U[c, :]

        if stop_if_too_bad:
            # usually we want to go this route for power flows
            if V.real.max() < 10:
                Scalc = cf.compute_power(Ybus, V)
                norm_f = cf.compute_fx_error(cf.compute_fx(Scalc, S0, no_slack, pq))
                converged = (norm_f <= tolerance) and (c % 2)  # we want an odd amount of coefficients

                if norm_f > norm_f_prev:
                    return U[c - 1, :], X[c - 1, :], Q[c - 1, :], V, iter_ - 1, converged

                norm_f_prev = norm_f
            else:
                # completely erroneous
                return U[c - 1, :], X[c - 1, :], Q[c - 1, :], V, iter_ - 1, converged
        else:
            # usually we want to go this route for sigma analysis
            Scalc = cf.compute_power(Ybus, V)
            norm_f = cf.compute_fx_error(cf.compute_fx(Scalc, S0, no_slack, pq))
            converged = (norm_f <= tolerance) and (c % 2)  # we want an odd amount of coefficients

        iter_ += 1
        c += 1

    return U, X, Q, V, iter_, converged


class HelmPreparation:
    """
    HelmPreparation
    """

    def __init__(self, sys_mat_factorization, Uini, Xini, Yslack, Vslack,
                 vec_P, vec_Q, Ysh, vec_W, pq, pv, pqpv, sl,
                 npqpv, nbus, pqpv_original, pq_original):
        """

        :param sys_mat_factorization:
        :param Uini:
        :param Xini:
        :param Yslack:
        :param Vslack:
        :param vec_P:
        :param vec_Q:
        :param Ysh:
        :param vec_W:
        :param pq:
        :param pv:
        :param pqpv:
        :param sl:
        :param npqpv:
        :param nbus:
        :param pqpv_original:
        """
        self.sys_mat_factorization = sys_mat_factorization
        self.Uini = Uini
        self.Xini = Xini
        self.Yslack = Yslack
        self.Vslack = Vslack
        self.vec_P = vec_P
        self.vec_Q = vec_Q
        self.Ysh = Ysh
        self.vec_W = vec_W
        self.pq = pq
        self.pv = pv
        self.sl = sl
        self.pqpv = pqpv
        self.npqpv = npqpv
        self.nbus = nbus
        self.pqpv_original = pqpv_original
        self.pq_original = pq_original


def helm_preparation_dY(Yseries, V0, S0, Ysh0, pq, pv, sl, pqpv, verbose: int = 0,
                        logger: Logger = None) -> HelmPreparation:
    """
    This function returns the constant objects to run many HELM simulations

    Based on the paper: Novel AC Distribution Factor for Efficient Outage Analysis
                        Rui Yao, Senior Member, IEEE, and Feng Qiu, Senior Member, IEEE

    :param Yseries: Admittance matrix of the series elements
    :param V0: vector of specified voltages
    :param S0: vector of specified power
    :param Ysh0: vector of shunt admittances (including the shunts of the Branches)
    :param pq: list of pq nodes
    :param pv: list of pv nodes
    :param sl: list of slack nodes
    :param pqpv: sorted list of pq and pv nodes
    :param verbose: print intermediate information
    :param logger: Logger object to store the debug info
    :return: U, X, Q, V, iterations
    """

    npqpv = len(pqpv)
    npv = len(pv)
    nsl = len(sl)
    nbus = Yseries.shape[0]

    # --------------------------- PREPARING IMPLEMENTATION -------------------------------------------------------------

    # build the reduced system
    Yred = Yseries[np.ix_(pqpv, pqpv)]  # admittance matrix without slack buses
    Yslack = -Yseries[np.ix_(pqpv, sl)]  # yes, it is the negative of this
    G = Yred.real.copy()  # real parts of Yij
    B = Yred.imag.copy()  # imaginary parts of Yij
    vec_P = S0.real[pqpv]
    vec_Q = S0.imag[pqpv]
    Vslack = V0[sl]
    Ysh = Ysh0[pqpv]
    Vm0 = np.abs(V0[pqpv])
    vec_W = Vm0 * Vm0

    # indices 0 based in the internal scheme
    nsl_counted = np.zeros(nbus, dtype=int)
    compt = 0
    for i in range(nbus):
        if i in sl:
            compt += 1
        nsl_counted[i] = compt

    pqpv_original = np.sort(np.r_[pq, pv])
    pq_original = np.sort(np.r_[pq])

    pq_ = pq - nsl_counted[pq]
    pv_ = pv - nsl_counted[pv]
    pqpv_ = np.sort(np.r_[pq_, pv_])

    # .......................CALCULATION OF TERMS [0] ------------------------------------------------------------------
    if nsl > 1:
        Uini = spsolve(Yred, Yslack.sum(axis=1))
    else:
        Uini = spsolve(Yred, Yslack)

    Xini = 1 / Uini

    # .......................CALCULATION OF THE MATRIX -----------------------------------------------------------------
    Upv = Uini[pv_]
    Xpv = Xini[pv_]
    VRE = coo_matrix((2 * Upv.real, (np.arange(npv), pv_)), shape=(npv, npqpv)).tocsc()
    VIM = coo_matrix((2 * Upv.imag, (np.arange(npv), pv_)), shape=(npv, npqpv)).tocsc()
    XIM = coo_matrix((-Xpv.imag, (pv_, np.arange(npv))), shape=(npqpv, npv)).tocsc()
    XRE = coo_matrix((Xpv.real, (pv_, np.arange(npv))), shape=(npqpv, npv)).tocsc()
    EMPTY = csc_matrix((npv, npv))

    MAT = vs((hs((G, -B, XIM)),
              hs((B, G, XRE)),
              hs((VRE, VIM, EMPTY))), format='csc')

    if verbose:
        logger.add_debug("MAT", MAT.toarray())

    # solve
    mat_factorized = factorized(MAT)

    return HelmPreparation(mat_factorized, Uini, Xini, Yslack, Vslack, vec_P, vec_Q, Ysh, vec_W,
                           pq_, pv_, pqpv_, sl, npqpv, nbus, pqpv_original, pq_original)


def helm_coefficients_dY(dY, sys_mat_factorization, Uini, Xini,
                         Yslack, Ysh, Ybus, vec_P, vec_Q, S0,
                         vec_W, V0, Vslack, pq, pv, pqpv, npqpv, nbus, sl, 
                         pqpv_original, pq_original, 
                         tolerance=1e-6, max_coeff=10):
    """
    Holomorphic Embedding LoadFlow Method as formulated by Josep Fanals Batllori in 2020
    This function just returns the coefficients for further usage in other routines

    Based on the paper: Novel AC Distribution Factor for Efficient Outage Analysis
                        Rui Yao, Senior Member, IEEE, and Feng Qiu, Senior Member, IEEE

    :param dY:
    :param sys_mat_factorization: factorized HELM system matrix
    :param Uini:
    :param Xini:
    :param Yslack:
    :param Ysh:
    :param Ybus:
    :param vec_P:
    :param vec_Q:
    :param S0: vector of specified power
    :param vec_W:
    :param V0: vector of specified voltages
    :param Vslack:
    :param pq: reduced scheme list of pq nodes
    :param pv: reduced scheme list of pv nodes
    :param pqpv: reduced scheme list of pq|pv nodes
    :param npqpv: number of pq and pv nodes
    :param nbus: number of nodes
    :param pqpv: sorted list of pq and pv nodes
    :param pq:list of pq nodes
    :param sl: list of slack nodes
    :param pqpv_original: sorted list of pq and pv nodes in the original order, considering slack
    :param pq_original: list of pq nodes in the original order, considering slack
    :param tolerance: target error (or tolerance)
    :param max_coeff: maximum number of coefficients
    :return: U, V, iter_, norm_f
    """

    AYred = dY[np.ix_(pqpv_original, pqpv_original)]  # difference admittance matrix without slack buses
    AYsl = dY[np.ix_(pqpv_original, sl)]  # difference admittance matrix involving slack buses

    # --------------------------- PREPARING IMPLEMENTATION -------------------------------------------------------------
    U = np.zeros((max_coeff + 1, npqpv), dtype=complex)  # voltages
    X = np.zeros((max_coeff + 1, npqpv), dtype=complex)  # compute X=1/conj(U)
    Q = np.zeros((max_coeff + 1, npqpv), dtype=complex)  # unknown reactive powers

    # .......................CALCULATION OF TERMS [0] ------------------------------------------------------------------
    U[0, :] = Uini
    X[0, :] = Xini

    # .......................CALCULATION OF TERMS [1] ------------------------------------------------------------------
    dval = np.zeros(npqpv, dtype=complex)

    # get the current Injections that appear due to the slack buses reduction
    I_inj_slack = Yslack[pqpv, :] * Vslack
    AI_inj_slack = AYsl[pqpv, :] * Vslack

    Ysl_sum = Yslack.sum(axis=1).A1
    AYsl_sum = AYsl.sum(axis=1).A1

    AIred = AYred @ U[0, :]

    dval[pq] = I_inj_slack[pq] + 1.0 * AI_inj_slack[pq] - Ysl_sum[pq] + (vec_P[pq] - vec_Q[pq] * 1j) * X[0, pq] - U[0, pq] * Ysh[pq] + 1.0 * AIred[pq]
    dval[pv] = I_inj_slack[pv] + 1.0 * AI_inj_slack[pv] - Ysl_sum[pv] + (vec_P[pv]) * X[0, pv] - U[0, pv] * Ysh[pv] + 1.0 * AIred[pv]

    # compose the right-hand side vector
    RHS = np.r_[dval.real, dval.imag, vec_W[pv] - (U[0, pv] * U[0, pv]).real]  # vec_W[pv_] - 1.0

    LHS = sys_mat_factorization(RHS)

    # update coefficients
    U[1, :] = LHS[:npqpv] + 1j * LHS[npqpv:2 * npqpv]
    Q[0, pv] = LHS[2 * npqpv:]
    X[1, :] = -X[0, :] * np.conj(U[1, :]) / np.conj(U[0, :])

    # .......................CALCULATION OF TERMS [>=2] ----------------------------------------------------------------
    iter_ = 1
    c = 2
    converged = False
    # V = np.empty(nbus, dtype=complex)
    # V[sl] = V0[sl]
    # V[pqpv] = U[:c, :].sum(axis=0)

    # Need the rebuild the voltage vector, as the passed pqpv indices are already reduced    
    V = np.zeros(nbus, dtype=complex)
    V[sl] = V0[sl]
    V[pqpv_original] = U[:c, :].sum(axis=0)

    norm_f = 0.0

    while c <= max_coeff and not converged:  # c defines the current depth

        AIred = AYred @ U[c - 1, :]

        dval[pq] = (vec_P[pq] - vec_Q[pq] * 1j) * X[c - 1, pq] - U[c - 1, pq] * Ysh[pq] + 1.0 * AIred[pq]
        dval[pv] = -1j * conv2(X, Q, c, pv) - U[c - 1, pv] * Ysh[pv] + X[c - 1, pv] * vec_P[pv] + 1.0 * AIred[pv]

        RHS = np.r_[dval.real, dval.imag, -conv3(U, U, c, pv).real]

        # LHS = spsolve(MAT, RHS)
        LHS = sys_mat_factorization(RHS)

        # update 
        U[c, :] = LHS[:npqpv] + 1j * LHS[npqpv:2 * npqpv]
        Q[c - 1, pv] = LHS[2 * npqpv:]
        X[c, :] = -conv1(U, X, c) / np.conj(U[0, :])

        # compute power mismatch
        V[pqpv_original] += U[c, :]

        if V.real.max() < 10:
            Scalc = cf.compute_power(Ybus - dY, V)
            norm_f = cf.compute_fx_error(cf.compute_fx(Scalc, S0, pqpv_original, pq_original))
            converged = (norm_f <= tolerance) and (c % 2)  # we want an odd amount of coefficients
        else:
            # completely erroneous
            break

        iter_ += 1
        c += 1

    return U, V, iter_, norm_f


def helm_josep(nc: NumericalCircuit,
               Ybus: CscMat, Yf: CscMat, Yt: CscMat, Yshunt_bus: CxVec,
               Yseries: CscMat, V0: CxVec, S0: CxVec, Ysh0: CxVec,
               pq: IntVec, pv: IntVec, vd: IntVec, no_slack: IntVec,
               tolerance: float = 1e-6, max_coefficients: int = 30, use_pade: bool = True,
               verbose: int = 0, logger: Logger = None) -> NumericPowerFlowResults:
    """
    Holomorphic Embedding LoadFlow Method as formulated by Josep Fanals Batllori in 2020
    :param nc: NumericalCircuit
    :param Ybus: Complete admittance matrix
    :param Yf: admittance from matrix
    :param Yt: admittance to matrix
    :param Yseries: Admittance matrix of the series elements
    :param V0: vector of specified voltages
    :param S0: vector of specified power
    :param Ysh0: vector of shunt admittances (including the shunt "legs" of the pi Branches)
    :param pq: list of pq nodes
    :param pv: list of pv nodes
    :param vd: list of slack nodes
    :param no_slack: sorted list of pq and pv nodes
    :param tolerance: target error (or tolerance)
    :param max_coefficients: maximum number of coefficients
    :param use_pade: Use the Padè approximation? otherwise, a simple summation is done
    :param verbose: print intermediate information
    :param logger: Logger object to store the debug info
    :return: V, converged, norm_f, Scalc, iter_, elapsed
    """

    start_time = time.time()

    if nc.bus_data.nbus < 2:
        return NumericPowerFlowResults(V=V0,
                                       Scalc=S0,
                                       m=np.ones(nc.nbr, dtype=float),
                                       tau=np.zeros(nc.nbr, dtype=float),
                                       Sf=np.zeros(nc.nbr, dtype=complex),
                                       St=np.zeros(nc.nbr, dtype=complex),
                                       If=np.zeros(nc.nbr, dtype=complex),
                                       It=np.zeros(nc.nbr, dtype=complex),
                                       loading=np.zeros(nc.nbr, dtype=complex),
                                       losses=np.zeros(nc.nbr, dtype=complex),
                                       Pf_vsc=np.zeros(nc.nvsc, dtype=float),
                                       St_vsc=np.zeros(nc.nvsc, dtype=complex),
                                       If_vsc=np.zeros(nc.nvsc, dtype=float),
                                       It_vsc=np.zeros(nc.nvsc, dtype=complex),
                                       losses_vsc=np.zeros(nc.nvsc, dtype=float),
                                       loading_vsc=np.zeros(nc.nvsc, dtype=float),
                                       Sf_hvdc=np.zeros(nc.nhvdc, dtype=complex),
                                       St_hvdc=np.zeros(nc.nhvdc, dtype=complex),
                                       losses_hvdc=np.zeros(nc.nhvdc, dtype=complex),
                                       loading_hvdc=np.zeros(nc.nhvdc, dtype=complex),
                                       norm_f=0.0,
                                       converged=False,
                                       iterations=0,
                                       elapsed=0.0)

    # compute the series of coefficients
    U, X, Q, V, iter_, converged = helm_coefficients_josep(Ybus=Ybus,
                                                           Yseries=Yseries,
                                                           V0=V0,
                                                           S0=S0,
                                                           Ysh0=Ysh0,
                                                           pq=pq,
                                                           pv=pv,
                                                           sl=vd,
                                                           no_slack=no_slack,
                                                           tolerance=tolerance,
                                                           max_coeff=max_coefficients,
                                                           verbose=bool(verbose),
                                                           stop_if_too_bad=True,
                                                           logger=logger)

    # --------------------------- RESULTS COMPOSITION ------------------------------------------------------------------
    if verbose:
        logger.add_debug('V coefficients\n', U)

    # compute the final voltage vector
    if use_pade:
        V = V0.copy()
        try:
            V[no_slack] = pade4all(max_coefficients - 1, U, 1)
        except RuntimeError:
            warn('Padè failed :(, using coefficients summation')
            V[no_slack] = U.sum(axis=0)

    # compute power mismatch
    Scalc = cf.compute_power(Ybus, V)
    norm_f = cf.compute_fx_error(cf.compute_fx(Scalc, S0, no_slack, pq))

    # check convergence
    converged = norm_f < tolerance

    elapsed = time.time() - start_time

    # Compute the Branches power and the slack buses power
    Sf, St, If, It, Vbranch, loading, losses, Sbus = cf.power_flow_post_process_nonlinear(
        Sbus=Scalc,
        V=V,
        F=nc.passive_branch_data.F,
        T=nc.passive_branch_data.T,
        pv=pv,
        vd=vd,
        Ybus=Ybus,
        Yf=Yf,
        Yt=Yt,
        Yshunt_bus=Yshunt_bus,
        branch_rates=nc.passive_branch_data.rates,
        Sbase=nc.Sbase
    )

    return NumericPowerFlowResults(
        V=V,
        Scalc=Scalc * nc.Sbase,
        m=np.ones(nc.nbr, dtype=float),
        tau=np.zeros(nc.nbr, dtype=float),
        Sf=Sf,
        St=St,
        If=If,
        It=It,
        loading=loading,
        losses=losses,
        Pf_vsc=np.zeros(nc.nvsc, dtype=float),
        St_vsc=np.zeros(nc.nvsc, dtype=complex),
        If_vsc=np.zeros(nc.nvsc, dtype=float),
        It_vsc=np.zeros(nc.nvsc, dtype=complex),
        losses_vsc=np.zeros(nc.nvsc, dtype=float),
        loading_vsc=np.zeros(nc.nvsc, dtype=float),
        Sf_hvdc=np.zeros(nc.nhvdc, dtype=complex),
        St_hvdc=np.zeros(nc.nhvdc, dtype=complex),
        losses_hvdc=np.zeros(nc.nhvdc, dtype=complex),
        loading_hvdc=np.zeros(nc.nhvdc, dtype=complex),
        norm_f=norm_f,
        converged=converged,
        iterations=iter_,
        elapsed=elapsed
    )
